feat: activate native async dispatch by pratyush618 · Pull Request #439 · ByteVeda/taskito

pratyush618 · 2026-07-17T03:30:34Z

Summary

Activates native async dispatch, dormant since inception because route flags were set on the returned TaskWrapper while the dispatch pool read the registry entry instead
Decouples worker shutdown from Python object lifetimes first, since the old drain relied on the async executor's result-sender pyclass being garbage-collected

Changes

Shutdown drain (Rust scheduler + Python): drain now exits on a refcount-independent condition — runtime thread finished AND Scheduler::in_flight_settled() — instead of waiting for the result channel to disconnect. Previously a retained exception (e.g. an on_failure hook collecting errors) pinned the frames reaching the executor's sender and stalled every shutdown for the full 30s timeout (measured 30.13s), past typical container grace periods. Gated off mesh mode since mesh-stolen jobs bypass in-flight tracking.
In-flight tracking ordering: track in-flight before the channel send, not after — a result handled between send and track previously leaked the slot forever.
Dispatch backpressure: channel-capacity pre-check added to both dispatch paths so a full job channel is skipped rather than claim -> Full -> rollback churning storage every tick.
Executor cleanup: AsyncTaskExecutor now drops its result sender on stop(); prefork no longer builds an executor it never uses, since holding that sender pin gave prefork the same 30s stall.
Activation: async route flags set on the registry entry and mirrored on TaskWrapper; max_in_flight widens to num_workers + async_concurrency only under the native-async feature, so sizing and activation land together. Non-native builds keep the flags inert with unchanged behavior.

Test plan

New regression test tests/worker/test_shutdown_drain.py — fails at 20s+ with the drain fix reverted, passes in ~1.4s with it; second test confirms coroutines run on the taskito-async-executor thread. Both skip on wheels built without native-async.
Full Python suite green on both builds: 1242 passed (native-async), 1240 passed (default, native tests skipped)
Rust: 196 tests default, 198 with native-async, all passing; postgres/redis feature checks green
Manual end-to-end: mixed sync/async workload with failing tasks, SIGTERM mid-flight — 0.21s shutdown, zero jobs left Running (40 completed, 5 dead-lettered)
Core scheduler unit tests confirmed to fail when the untrack-ordering or channel pre-check fixes are individually reverted

Summary by CodeRabbit

Bug Fixes
- Improved dispatch when worker capacity is unavailable, preventing stale scheduling and unnecessary job claiming.
- Fixed in-flight accounting so failures during handoff or closed channels no longer leak capacity.
- Enhanced bounded shutdown/draining so workers exit promptly once work is complete, including during failures.
- Improved native-async task routing consistency and ensured shutdown cleanup releases pinned resources.
New Features
- Added a public signal to check whether bounded in-flight work has fully settled.
Tests
- Expanded coverage for dispatch rollback/capacity, bounded settled behavior, native-async execution, and fast shutdown.

A result handled between try_send and track_in_flight released nothing, leaking the slot for the life of the worker.

Every dequeue against a full channel claimed and then rolled back on try_send, churning storage for nothing.

Waiting for the result channel to disconnect ties shutdown to the async result sender's Python refcount: a failed task's traceback pins it, burning the full drain timeout on every shutdown.

A pinned frame can keep the executor alive indefinitely, so the sender release cannot wait for garbage collection.

Route flags now live on the registry entry the pool reads, and the in-flight cap widens to num_workers + async_concurrency since the two budgets are real once coroutines run on the executor.

coderabbitai · 2026-07-17T03:31:03Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 2181c8b8-8113-47f8-8feb-9d75d992f9fa

📥 Commits

Reviewing files that changed from the base of the PR and between b457b48 and 4f0f98b.

📒 Files selected for processing (1)

sdks/python/tests/worker/test_native_async.py

🚧 Files skipped from review as they are similar to previous changes (1)

sdks/python/tests/worker/test_native_async.py

📝 Walkthrough

Walkthrough

The scheduler now avoids full channels, rolls back failed handoffs, and reports bounded in-flight settlement. Native-async workers account for async concurrency and stop cleanly after settled work. Async markers and executor sender cleanup are covered by new tests.

Changes

Dispatch and shutdown lifecycle

Layer / File(s)	Summary
Scheduler capacity guards and rollback `crates/taskito-core/src/scheduler/mod.rs`, `crates/taskito-core/src/scheduler/poller.rs`	Dispatch skips unavailable capacity, limits batch dequeue size, untracks failed handoffs, and exposes `in_flight_settled()`.
Worker settlement and executor shutdown `crates/taskito-python/src/py_queue/worker.rs`	Native-async capacity contributes to `max_in_flight`; shutdown exits when runtime and scheduler work settle, then stops the retained executor.
Async dispatch markers and sender release `sdks/python/taskito/mixins/decorators.py`, `sdks/python/taskito/async_support/executor.py`	Async markers are assigned to the internal callable, and `stop()` releases the result sender after executor shutdown.
Lifecycle regression coverage `crates/taskito-core/src/scheduler/mod.rs`, `sdks/python/tests/worker/test_native_async.py`, `sdks/python/tests/worker/test_shutdown_drain.py`	Tests cover rollback, settlement reporting, sender release, native executor execution, worker termination, and prompt shutdown after task failure.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Queue
  participant run_worker
  participant Scheduler
  participant AsyncTaskExecutor
  Queue->>run_worker: shutdown()
  run_worker->>Scheduler: Check in_flight_settled()
  run_worker->>run_worker: Check runtime_handle.is_finished()
  run_worker->>AsyncTaskExecutor: stop()
  AsyncTaskExecutor-->>run_worker: Event loop stopped and sender released

Possibly related PRs

ByteVeda/taskito#433: Related scheduler in-flight tracking and dispatch rollback changes.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: activating native async dispatch and related lifecycle updates.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch feat/native-async-activation

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 2

🧹 Nitpick comments (1)

crates/taskito-core/src/scheduler/mod.rs (1)
1212-1237: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Exercise the actual TrySendError::Full rollback path.

The pre-filled channel makes tick_dispatch() return at the new capacity check, so this test never tracks or rolls back a job. It duplicates the skip test at Lines 1588-1622 and leaves the check-to-send race untested.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/taskito-core/src/scheduler/mod.rs` around lines 1212 - 1237, The test
test_full_channel_rollback_does_not_leak_in_flight_slot currently exits at the
capacity check instead of exercising rollback after tracking. Arrange the
scheduler and channel state so capacity is available when tick_dispatch begins,
then fill the channel between the check and try_send to force
TrySendError::Full; retain assertions that dispatch fails and
in_flight_remaining() returns Some(8), while avoiding duplication of the
capacity-skip test.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@crates/taskito-python/src/py_queue/worker.rs`:
- Around line 288-297: Update the max_in_flight calculation in the worker
capacity logic so prefork mode always uses only self.num_workers, including
native-async builds; apply the same correction to the duplicated logic noted in
the comment. Retain self.num_workers + async_concurrency only for pool modes
that actually create and use the native async executor.

In `@sdks/python/taskito/async_support/executor.py`:
- Around line 270-283: Update stop() to check whether self._loop is not closed
rather than whether it is currently running before calling
call_soon_threadsafe(self._loop.stop). Keep the existing thread join, timeout
warning, and sender release behavior unchanged.

---

Nitpick comments:
In `@crates/taskito-core/src/scheduler/mod.rs`:
- Around line 1212-1237: The test
test_full_channel_rollback_does_not_leak_in_flight_slot currently exits at the
capacity check instead of exercising rollback after tracking. Arrange the
scheduler and channel state so capacity is available when tick_dispatch begins,
then fill the channel between the check and try_send to force
TrySendError::Full; retain assertions that dispatch fails and
in_flight_remaining() returns Some(8), while avoiding duplication of the
capacity-skip test.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: f89d5588-c6ce-4358-9476-090d49146a47

📥 Commits

Reviewing files that changed from the base of the PR and between 33867c6 and a60a7d2.

📒 Files selected for processing (7)

crates/taskito-core/src/scheduler/mod.rs
crates/taskito-core/src/scheduler/poller.rs
crates/taskito-python/src/py_queue/worker.rs
sdks/python/taskito/async_support/executor.py
sdks/python/taskito/mixins/decorators.py
sdks/python/tests/worker/test_native_async.py
sdks/python/tests/worker/test_shutdown_drain.py

is_running() is false until the thread enters run_forever, so a stop racing start skipped loop.stop and leaked the thread.

Prefork builds no async executor, so the widened native-async cap promised capacity nothing could execute.

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@sdks/python/tests/worker/test_native_async.py`:
- Line 7: Update the shutdown-race test around the loop-stop handoff to
synchronize with an explicit event that is set when the loop has been inspected
or stopped, wait for that event, and only then release entered. Remove the
time.sleep-based delay and the now-unused time import, preserving the test’s
validation of stop() behavior.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 3cf16bce-ca94-40e2-a084-ec25371d33a1

📥 Commits

Reviewing files that changed from the base of the PR and between a60a7d2 and b457b48.

📒 Files selected for processing (3)

crates/taskito-python/src/py_queue/worker.rs
sdks/python/taskito/async_support/executor.py
sdks/python/tests/worker/test_native_async.py

🚧 Files skipped from review as they are similar to previous changes (2)

sdks/python/taskito/async_support/executor.py
crates/taskito-python/src/py_queue/worker.rs

Signal from the join stop() reaches after its loop check instead of sleeping and hoping the check already ran.

The old bound of 8 assumed the in-flight cap binds; on native-async builds sync work binds on the channel plus finalize lag (CI hit 9).

pratyush618 added 5 commits July 17, 2026 08:17

fix(scaling): track in-flight before channel send

f65193d

A result handled between try_send and track_in_flight released nothing, leaking the slot for the life of the worker.

perf(scaling): skip dispatch while job channel is full

ae8768b

Every dequeue against a full channel claimed and then rolled back on try_send, churning storage for nothing.

fix(python): end shutdown drain when in-flight work settles

b3b7d25

Waiting for the result channel to disconnect ties shutdown to the async result sender's Python refcount: a failed task's traceback pins it, burning the full drain timeout on every shutdown.

fix(python): release async result sender on executor stop

bca5dfe

A pinned frame can keep the executor alive indefinitely, so the sender release cannot wait for garbage collection.

feat(python): activate native async dispatch

a60a7d2

Route flags now live on the registry entry the pool reads, and the in-flight cap widens to num_workers + async_concurrency since the two budgets are real once coroutines run on the executor.

github-actions Bot added rust python scheduler tests labels Jul 17, 2026

coderabbitai Bot reviewed Jul 17, 2026

View reviewed changes

Comment thread crates/taskito-python/src/py_queue/worker.rs

Comment thread sdks/python/taskito/async_support/executor.py

pratyush618 added 2 commits July 17, 2026 09:25

fix(python): stop executor loop even before it starts running

7785957

is_running() is false until the thread enters run_forever, so a stop racing start skipped loop.stop and leaked the thread.

fix(python): keep prefork in-flight cap at num_workers

b457b48

Prefork builds no async executor, so the widened native-async cap promised capacity nothing could execute.

coderabbitai Bot reviewed Jul 17, 2026

View reviewed changes

Comment thread sdks/python/tests/worker/test_native_async.py Outdated

pratyush618 added 2 commits July 17, 2026 09:32

test(python): make executor stop-race test deterministic

4f0f98b

Signal from the join stop() reaches after its loop check instead of sleeping and hoping the check already ran.

test(python): derive in-flight bound from worker pipeline

8946707

The old bound of 8 assumed the in-flight cap binds; on native-async builds sync work binds on the channel plus finalize lag (CI hit 9).

pratyush618 merged commit 925bf61 into master Jul 17, 2026
34 checks passed

pratyush618 deleted the feat/native-async-activation branch July 17, 2026 05:41

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: activate native async dispatch#439

feat: activate native async dispatch#439
pratyush618 merged 9 commits into
masterfrom
feat/native-async-activation

pratyush618 commented Jul 17, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Jul 17, 2026 •

edited

Loading

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

pratyush618 commented Jul 17, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Test plan

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

pratyush618 commented Jul 17, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 17, 2026 •

edited

Loading